#!/usr/bin/env python
'''
This program compares series expansion made by Rajiv Singh and data from simulations
'''
from __future__ import division

import os
import sys
import matplotlib.pyplot

sys.path += ['..']

import common
import re
import math
import scipy
from sympy import *

path = os.path.join('..', '..', 'results')


modelname = common.choice(path)
J = common.choice(os.path.join(path, modelname))
B = common.choice(os.path.join(path, modelname, J))
size = common.choice( os.path.join(path, modelname, J, B) )

if type(size) == str:
    size = [size]

dT = common.choice( os.path.join(path, modelname, J, B, size[0]) )

variable = common.choice(os.path.join(path, modelname, J, B, size[0]))

def Entropy(T, E):
    '''
    returns entropy
    '''
    beta = [1 / tt for tt in T]
    S = []
    for i in xrange(len(E)):
        S += [ math.log(2) + E[i] / T[i] + scipy.trapz(E[i:], beta[i:]) ]
        
    minS = S[0]
    return [ts - minS for ts in S]  


for tsize in size:
    result_path_J = os.path.join( path, modelname, J, B, tsize, dT, 'Energy_J' )
    result_J = []
    
    for T in os.listdir(result_path_J):
        E_J = eval( open( os.path.join(result_path_J, T)  ).read() )
        result_J += [ ( float(T.replace('T','')), E_J ) ]
    
    result_J.sort()
    
    x_J = [ tx[0] for tx in result_J ]
    y_J = [ tx[1] for tx in result_J ]
    
    result_path_B = os.path.join( path, modelname, J, B, tsize, dT, 'Energy_B' )
    result_B = []
    
    for T in os.listdir(result_path_B):
        E_B = eval( open( os.path.join(result_path_B, T)  ).read() )
        result_B += [ ( float(T.replace('T','')), E_B ) ]
    
    result_B.sort()
    
    x_B = [ tx[0] for tx in result_B ]
    y_B = [ tx[1] for tx in result_B ]
    
    y_J_B = []
    
    for i in xrange(len(y_J)):
        y_J_B += [y_J[i] + y_B[i]]
        
    matplotlib.pyplot.plot(x_J, y_J_B, "*", label = '$N = %s$'%(tsize.replace('N','')))




B = eval(B.replace('B_',''))
#print B
#eRange = [ Energy(1, tt, B) for tt in x_J]



matplotlib.pyplot.xlabel('Temperature')
matplotlib.pyplot.ylabel('Energy')
matplotlib.pyplot.title('Comparison Energy from simulution vs Energy from fancy series expansion\n for $B = %s$'%(str(B)))

#And now I'm working with series expansion

var('h J beta')

Z1 = 2 * cosh(beta * h)
Z2 = 2 * cosh(beta * sqrt(4 * h**2 + J**2) ) + 2 * cosh(beta * J)



E1 = - log(Z1).diff(beta) - 3 * log(Z2).diff(beta) + 6 * log(Z1).diff(beta)

E1 = E1.subs(J, 1).subs(h, B)

tErange1 = [ E1.subs(beta, 1 / tt) for tt in x_J  ]

matplotlib.pyplot.plot(x_J, tErange1, "*", label = 'series expansion $J / B \ll 1$')


E2 = - h * tanh(beta * h) - 3 / 4 * J**2 * beta / cosh(beta *h)**2* ( 2 + 1 / (2 * beta * h) * sinh(2 * beta * h) + cosh(2 * beta * h)) + 3 / 2 * (beta * J)**2 * h / cosh(beta * h)**2 * tanh(beta * h) * (1 + 1 / (2 * beta * h) * sinh(2 * beta * h))

#print latex(E2)

E2 = E2.subs(J, 1).subs(h, B)


tErange2 = [ E2.subs(beta, 1 / tt) for tt in x_J  ]

matplotlib.pyplot.plot(x_J, tErange2, "^", label = 'series expansion $J / B \ll 1$, $J / T \ll 1$')
matplotlib.pyplot.legend(loc = 'upper left')


matplotlib.pyplot.show()